from PyQt5.QtGui import QStandardItemModel, QStandardItem, QBrush, QColor
from PyQt5.QtCore import Qt
import json


def load_data_from_json(file_path):
    with open(file_path, 'r', encoding='utf-8') as file:
        data = json.load(file)
    return data

class SearchableModel(QStandardItemModel):
    def __init__(self, data):
        super().__init__()
        self.original_data = data
        self.create_model(data)

    def create_model(self, data):
        self.invisibleRootItem().removeRows(0, self.rowCount())
        units = {}

        for unit_name, unit_data in data.items():
            units[unit_name] = QStandardItem(unit_name)
            units[unit_name].setData(unit_data.get('business', ''), Qt.UserRole + 1)

        for unit_name, unit_data in data.items():
            superior_unit = unit_data['superior_unit']
            if superior_unit == "无上级单位":
                self.invisibleRootItem().appendRow(units[unit_name])
            else:
                if superior_unit in units:
                    units[superior_unit].appendRow(units[unit_name])

    def search(self, text, column=0):
        self.invisibleRootItem().removeRows(0, self.rowCount())

        if not text:
            self.create_model(self.original_data)
            self.layoutChanged.emit()
            return

        filtered_data = {}
        for key, value in self.original_data.items():
            if text.lower() in key.lower():
                filtered_data[key] = value
                superior_unit = value['superior_unit']
                while superior_unit != "无上级单位" and superior_unit not in filtered_data:
                    filtered_data[superior_unit] = self.original_data[superior_unit]
                    superior_unit = self.original_data[superior_unit]['superior_unit']

        units = {}
        for unit_name, unit_data in filtered_data.items():
            units[unit_name] = QStandardItem(unit_name)
            units[unit_name].setData(unit_data.get('business', ''), Qt.UserRole + 1)
            if text.lower() in unit_name.lower():
                units[unit_name].setData(QBrush(QColor(255, 255, 0)), Qt.BackgroundRole)  # 设置背景色为黄色

        for unit_name, unit_data in filtered_data.items():
            superior_unit = unit_data['superior_unit']
            if superior_unit == "无上级单位":
                self.invisibleRootItem().appendRow(units[unit_name])
            else:
                if superior_unit in units:
                    units[superior_unit].appendRow(units[unit_name])

        self.layoutChanged.emit()
    def expand_item(self, item):
        if item.hasChildren():
            item.setExpanded(True)
            for i in range(item.rowCount()):
                self.expand_item(item.child(i))

def create_tree_model(data):
    return SearchableModel(data)